import os

import pandas as pd
import pymssql
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
class MyVanna(OpenAI_Chat):
    def __init__(self, client, config=None):  
        super().__init__(client=client, config=config)
        self.conn = None
        
    def connect_to_mssql(self, server=None, database=None, user=None, password=None):
        """使用 pymssql 替代默认的 pyodbc 连接"""
        try:
            self.conn = pymssql.connect(
                server=server or os.getenv('DB_SERVER'),
                database=database or os.getenv('DB_NAME'),
                user=user or os.getenv('DB_USER'),
                password=password or os.getenv('DB_PASSWORD')
            )
            return True
        except Exception as e:
            print(f"数据库连接错误: {str(e)}")
            return False
            
    def run_sql(self, sql: str) -> pd.DataFrame:
        """重写 run_sql 方法，使用 pymssql 执行 SQL"""
        if self.conn is None:
            raise Exception("请先连接到数据库")
        
        try:
            return pd.read_sql(sql, self.conn)
        except Exception as e:
            raise Exception(f"SQL执行错误: {str(e)}")